from torch_geometric.data import HeteroData
import torch_geometric.transforms  as T
import torch
data = HeteroData()

data['paper'].x = torch.randn([6,4])
data['author'].x = torch.randn([3,4])
data['institution'].x = torch.randn([2,4])
data['field_of_study'].x = torch.randn([2,4])

data['paper', 'cites', 'paper'].edge_index = torch.LongTensor([[1,3,5],[2,2,4]])
data['author', 'writes', 'paper'].edge_index = torch.LongTensor([[0,0,1,1,1,2],[0,1,2,3,4,5]])
data['author', 'affiliated_with', 'institution'].edge_index = torch.LongTensor([[0,1,2],[0,1,1]])
data['paper', 'has_topic', 'field_of_study'].edge_index =  torch.LongTensor([[0,1,2,3,4,5],[0,0,1,1,1,1]])


# data['paper', 'cites', 'paper'].edge_attr = ... # [num_edges_cites, num_features_cites]
# data['author', 'writes', 'paper'].edge_attr = ... # [num_edges_writes, num_features_writes]
# data['author', 'affiliated_with', 'institution'].edge_attr = ... # [num_edges_affiliated, num_features_affiliated]
# data['paper', 'has_topic', 'field_of_study'].edge_attr = ... # [num_edges_topic, num_features_topic]
del data['field_of_study']  # Deleting 'field_of_study' node type
del data['has_topic']  # Deleting 'has_topic' edge type

print(data["paper"].x)
data = T.ToUndirected()(data)
data = T.AddSelfLoops()(data)
data = T.NormalizeFeatures()(data)
print(data.edge_index_dict)
"""
{('paper', 'cites', 'paper'): 
增加自环 AddSelfLoops 中指的是在同构点之间增加自环，在异构点的边的两端点是不加自环的。
tensor([[1, 2, 2, 3, 4, 5, 0, 1, 2, 3, 4, 5],
        [2, 1, 3, 2, 5, 4, 0, 1, 2, 3, 4, 5]]),
('author', 'writes', 'paper'): 
tensor([[0, 0, 1, 1, 1, 2],
        [0, 1, 2, 3, 4, 5]]), 
('author', 'affiliated_with', 'institution'): 
tensor([[0, 1, 2],
        [0, 1, 1]]),
        
增加双向 ToUndirected 在同构边上增加双倍的指向，创建反向的异构边。
('paper', 'rev_writes', 'author'): 
tensor([[0, 1, 2, 3, 4, 5],
        [0, 0, 1, 1, 1, 2]]), 
('institution', 'rev_affiliated_with', 'author'): 
tensor([[0, 1, 1],
        [0, 1, 2]])}
"""
print(data["paper"].x)
print("##################################################")
print(data)

"""
HeteroData(
  paper={ x=[6, 4] },
  author={ x=[3, 4] },
  institution={ x=[2, 4] },
  (paper, cites, paper)={ edge_index=[2, 12] },
  (author, writes, paper)={ edge_index=[2, 6] },
  (author, affiliated_with, institution)={ edge_index=[2, 3] },
  (paper, rev_writes, author)={ edge_index=[2, 6] },
  (institution, rev_affiliated_with, author)={ edge_index=[2, 3] }
)

"""
for store in data.stores:
    # print(store.items)
    print(store.items(*["edge_index"]),"#############")
    for key, value in store.items(*["x"]):
        print(key,"####",value)
    # print(store,store.items)

"""
with torch.no_grad():  # Initialize lazy modules.
    out = model(data.x_dict, data.edge_index_dict)
"""
for edge_type in data.edge_types:
    src, _, dst = edge_type
    assert data[edge_type].edge_index[0].max() < data[src].num_nodes
    assert data[edge_type].edge_index[1].max() < data[dst].num_nodes
